"""Combine testing results of the three models to get final accuracy."""
from __future__ import division
from __future__ import print_function
import argparse
import numpy as np


def main():
    parser = argparse.ArgumentParser(description="combine predictions")
    parser.add_argument('--iframe', type=str, required=True,
                        help='iframe score file.')
    parser.add_argument('--mv', type=str, required=True,
                        help='motion vector score file.')
    parser.add_argument('--res', type=str, required=True,
                        help='residual score file.')

    parser.add_argument('--wi', type=float, default=2.0,
                        help='iframe weight.')
    parser.add_argument('--wm', type=float, default=1.0,
                        help='motion vector weight.')
    parser.add_argument('--wr', type=float, default=1.0,
                        help='residual weight.')

    args = parser.parse_args()

    with np.load(args.iframe) as iframe:
        with np.load(args.mv) as mv:
            with np.load(args.res) as residual:
                n = len(mv['names'])

                i_score = np.array([score[0][0] for score in iframe['scores']])
                mv_score = np.array([score[0][0] for score in mv['scores']])
                res_score = np.array([score[0][0] for score in residual['scores']])
                # print(i_score.shape)

                i_label = np.array([score[1] for score in iframe['scores']])
                mv_label = np.array([score[1] for score in mv['scores']])
                res_label = np.array([score[1] for score in residual['scores']])
                assert np.alltrue(i_label == mv_label) and np.alltrue(i_label == res_label)

                print('I-frame')
                ODRup, ODRdown = 0, 0
                for clas in range(i_score.shape[1]):
                    i_pre = np.argmax(i_score, axis=1)
                    up = np.count_nonzero((i_pre == i_label) &
                                          (i_label == clas))
                    down = np.count_nonzero(i_label == clas)
                    if clas != 0:
                        ODRup += up
                        ODRdown += down
                    print('\tEv' + str(clas+100), 'accu:', up/down)
                print('\tODR', ODRup/ODRdown)

                print('mv')
                ODRup, ODRdown = 0, 0
                for clas in range(i_score.shape[1]):
                    mv_pre = np.argmax(mv_score, axis=1)
                    up = np.count_nonzero((mv_pre == mv_label) &
                                          (mv_label == clas))
                    down = np.count_nonzero(mv_label == clas)
                    if clas != 0:
                        ODRup += up
                        ODRdown += down
                    print('\tEv' + str(clas+100), 'accu:', up/down)
                print('\tODR', ODRup/ODRdown)

                print('res')
                ODRup, ODRdown = 0, 0
                for clas in range(i_score.shape[1]):
                    res_pre = np.argmax(res_score, axis=1)
                    up = np.count_nonzero((res_pre == res_label) &
                                          (res_label == clas))
                    down = np.count_nonzero(res_label == clas)
                    if clas != 0:
                        ODRup += up
                        ODRdown += down
                    print('\tEv' + str(clas+100), 'accu:', up/down)
                print('\tODR', ODRup/ODRdown)


                combined_score = i_score * args.wi + mv_score * args.wm + res_score * args.wr

                print('combine')
                ODRup, ODRdown = 0, 0
                for clas in range(i_score.shape[1]):
                    com_pre = np.argmax(combined_score, axis=1)
                    up = np.count_nonzero((com_pre == i_label) &
                                          (i_label == clas))
                    down = np.count_nonzero(i_label == clas)
                    if clas != 0:
                        ODRup += up
                        ODRdown += down
                    print('\tEv' + str(clas+100), 'accu:', up/down)
                print('\tODR', ODRup/ODRdown)

                accuracy = float(sum(np.argmax(combined_score, axis=1) == i_label)) / n
                print('Accuracy: %f (%d).' % (accuracy, n))

if __name__ == '__main__':
    main()
